import json
import os
from abc import ABC
from copy import deepcopy

from dill.pointers import children
from markdown_it.common.html_re import attribute
import collections


class BaseNode(object):
    def __init__(self) -> None:
        super().__init__()

class CodeSnippet(BaseNode):
    def __init__(self, description, codes):
        self.description = description
        self.codes = codes

    def __repr__(self):
        return f"<CodeSnippet: {self.description}>"
    
    def to_string(self):
        return f"description: {self.description} \n codes: {self.codes}"

class ImagePaths(BaseNode):
    def __init__(self, paths):
        self.paths = paths  # dict格式 {图片编号: 路径}

    def __repr__(self):
        return f"<ImagePaths: {len(self.paths)} images>"

class DepthDict(dict):
    def __init__(self,current_level_item,depth):
        super().__init__({current_level_item.id:current_level_item})
        self.depth_table = {depth:[{current_level_item.id:current_level_item}]}

    def update_with_depth(self,new_dict):
        self.depth_table.update(new_dict.depth_table)
    def append_with_depth(self,new_dict):
        self.update(new_dict)
        for depth,id in new_dict.depth_table.items():
            dict_item = self.depth_table.get(depth)
            if dict_item is None:
                self.depth_table[depth] = [new_dict]
            else:
                self.depth_table[depth].append(new_dict)
    # def update
    @property            
    def depth_table_size(self):
        return len(self.depth_table)

    def get_items_with_depth(self,depth):
        return self.depth_table[depth]

class HierarchyObjects(BaseNode):
    def __init__(self, children: dict,id,depth:int):
        self.id:str = id
        self.children = children  # dict格式 {子章节编号: Section}
        self.all_item_with_depth = DepthDict(self,depth)
        self.depth = depth
        self.get_all_item_with_depth()
    
    def get_all_item_with_depth(self):
        if len(self.children) == 0:
            return
        # self.all_item_with_depth.update_with_depth(self.children,self.depth+1)
        for key,value in self.children.items():
            self.all_item_with_depth.append_with_depth(value.all_item_with_depth)

    def get_items_with_depth(self,start_relative_depth = 0):
        if start_relative_depth > self.all_item_with_depth.depth_table_size:
            # too depth error
            raise ValueError(f"depth must less than {self.all_item_with_depth.depth_table_size}, but got {start_relative_depth}")
        all_depth = range(start_relative_depth+self.depth,self.all_item_with_depth.depth_table_size)
        ret = []
        print("all_depth==",all_depth)
        for i in all_depth:
            ret.append(self.all_item_with_depth.get_items_with_depth(i))
        return ret

class Book(HierarchyObjects):
    def __init__(self, children: dict = {},id = "root", depth=0):
        super().__init__(children,"root",-1)
    def __getitem__(self,key):
        return self.children[key]
    def get(self,key,default={}):
        return  self.children.get(key,default)
    def update(self,other_book:"Book"):
        self.children.update(other_book.children)
    
class Section(HierarchyObjects):
    def __init__(self,id = -1, title="", summary="", original_content="", code_snippets={}, image_paths=ImagePaths({}), children={}, depth=0):
        super().__init__(children,id,depth)
        self.title: str = title
        self.summary: str = summary
        self.original_content: str = original_content
        self.code_snippets:dict[str,CodeSnippet] = code_snippets  # dict格式 {代码段编号: CodeSnippet}
        self.original_content_with_codes: str = self.replace_code_snippets_in_content()
        self.image_paths = image_paths

    def __repr__(self):
        return f"<Section {self.id}: {self.title}>"

    def replace_code_snippets_in_content(self):
        tmp_content = deepcopy(self.original_content)
        for code_idx in self.code_snippets.keys():
            code_snippet_name = f"【代码段{code_idx}】"
            tmp_content = tmp_content.replace(code_snippet_name,f"{self.code_snippets[code_idx].to_string()}")
        return tmp_content

    def to_prompt(self):
        return f"Section id: {self.id} \n Section title: {self.title} \n Section summary: {self.summary} \n Original content: {self.original_content_with_codes}"

    def get_all_code_snippets(self):
        ret = ""
        for code_idx in self.code_snippets.keys():
            ret += self.code_snippets[code_idx].to_string() + "\n"
        return ret
    @classmethod
    def from_book(cls, book:"Book"):
        return cls(children = book.children)

class JsonParser:
    global_code_snippets = {}
    global_image_path = {}
    
    @staticmethod
    def _get_code_snippets_from_global(section_original_content:str):
        if len(JsonParser.global_code_snippets) == 0:
            return {}
        ret = {}
        for key in JsonParser.global_code_snippets.keys():
            code_snippets_id = f"【代码段{key}】"
            if (section_original_content.find(code_snippets_id) != -1):
                code = JsonParser.global_code_snippets[key]
                ret[key] = code
        return ret

    @staticmethod
    def _get_image_path_from_global(section_original_content:str):
        if len(JsonParser.global_image_path) == 0:
            return {}
        ret = {}
        for key in JsonParser.global_image_path.keys():
            code_snippets_id = f"【图片{key}】"
            if (section_original_content.find(code_snippets_id) != -1):
                code = JsonParser.global_image_path[key]
                ret[key] = code
        return ret

    @staticmethod
    def parse_json(json_data):
        """JSON解析入口函数"""
        chapters = {}
        if json_data.get('代码段') is not None:
            JsonParser.global_code_snippets = json_data.get('代码段',{})
        if json_data.get('图片路径') is not None:
            JsonParser.global_image_path = json_data.get('图片路径',{})

        for id, info in json_data.items():
            if info.get('代码段') is not None:
                JsonParser.global_code_snippets.update(info.get('代码段',{}))
            if info.get('图片路径') is not None:
                JsonParser.global_image_path.update(info.get('图片路径',{}))

        for chapter_id, chapter_info in json_data.items():
            if chapter_id == '代码段' or chapter_id == '图片路径':
                continue
            chapters[chapter_id] = JsonParser._build_section(chapter_id,chapter_info,depth = 0)
        return Book(chapters)

    @staticmethod
    def _build_children(children_dict,depth = 0):
        """构建子章节结构"""
        children = {}
        for section_id, section_info in children_dict.items():
            children[section_id] = JsonParser._build_section(section_id,section_info, depth+1)
        return children

    @staticmethod
    def _build_section(section_id,section_dict, depth):
        """构建子章节内容"""
        original_content = section_dict.get('原文内容', '')
        code_snippets_dict = section_dict.get('代码段', {})
        code_snippets_dict.update(JsonParser._get_code_snippets_from_global(original_content))
        image_paths_dict = section_dict.get('图片路径', {})
        image_paths_dict.update(JsonParser._get_image_path_from_global(original_content))
        code_snippets = JsonParser._build_code_snippets(code_snippets_dict)
        children = JsonParser._build_children(section_dict.get('children',{}),depth)
        return Section(
            id = section_id,
            title=section_dict['title'],
            summary=section_dict['summary'],
            original_content=original_content,
            code_snippets=code_snippets,
            image_paths=ImagePaths(image_paths_dict),
            children = children,
            depth=depth,
        )

    @staticmethod
    def _build_code_snippets(snippets_dict):
        """构建代码片段集合"""
        return {
            snippet_id: CodeSnippet(
                description=info['description'],
                codes=info['codes']
            )
            for snippet_id, info in snippets_dict.items()
        }

def read_json_files(folder_path):
    all_json_strs = []
    if not os.path.isdir(folder_path):
        raise  FileExistsError(f"Error: {folder_path} is not a valid directory.")
    
    for filename in os.listdir(folder_path):
        if filename.endswith(".json"):
            file_path = os.path.join(folder_path, filename)
            print(f"Reading file: {filename}")
            # try:
            with open(file_path, "r", encoding="utf-8") as file:
                json_data = json.load(file)
                current_chapter = JsonParser.parse_json(json_data)
                all_json_strs.append(current_chapter)
                    # print(json.dumps(data, indent=4, ensure_ascii=False))
            # except Exception as e:
            #     print(f"Failed to read {filename}: {e},{e.with_traceback()}")
            # print("-----------------------------")
    return all_json_strs

def load_jsons_from_path(path):
    all_chapters = Book()
    if os.path.isdir(path):
        for file_name in os.listdir(path):
            if file_name.endswith((".json")):  # 只读取 Markdown 文件
                file_path = os.path.join(path, file_name)
                with open(file_path, "r", encoding="utf-8") as file:
                    json_data = json.load(file)
                    current_chapter = JsonParser.parse_json(json_data)
                    all_chapters.update(current_chapter)
    return all_chapters

def load_jsons_from_dir(path):
    json_strs = read_json_files(path)
    all_chapters = []
    for json_data in json_strs:
            current_chapter = JsonParser.parse_json(json_data)
            all_chapters.append(current_chapter)
    return all_chapters    


def load_and_add_json(paths):
    all_chapters = Book()
    for p in paths:
        with open(p, encoding='utf-8') as f:
            json_data = json.load(f)
            current_chapter = JsonParser.parse_json(json_data)
            all_chapters.update(current_chapter)
    return all_chapters

def load_jsons(paths):
    all_chapters = []
    for p in paths:
        with open(p, encoding='utf-8') as f:
            json_data = json.load(f)
            current_chapter = JsonParser.parse_json(json_data)
            all_chapters.append(current_chapter)
    return all_chapters

if __name__ == "__main__":
    # 从文件加载JSON数据
    with open('/home/ma-user/work/users/wty/kernel-gen/data/chapter_contents/16.json', encoding='utf-8') as f:
        json_data = json.load(f)
    
    # 动态构建对象结构
    chapters = JsonParser.parse_json(json_data)
    
    # 验证结构
    chapter3 = chapters['6']
    # print(chapter3.children['6.2'].summary)
    print(chapter3.children['6.2'].children['6.2.1'].children['6.2.1.1'].code_snippets['3'].codes)
    print("======================")
    print(chapter3.all_item_with_depth['6.2.1.1'].code_snippets['3'].codes)
    print("======================")
    print(chapter3.all_item_with_depth['6.2.1.2'].depth)
    print("======================")
    print(chapter3.all_item_with_depth['6.2.1.1'].all_item_with_depth)
    # print(chapter3.children['15.1'].code_snippets)